import com.feriki.msync.config.ComplexConnectConfig;
import com.feriki.msync.pool.MilvusConnectionPool;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResults;
import io.milvus.param.R;
import io.milvus.param.dml.SearchParam;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.tuple.Pair;

import java.io.File;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;

/**
 * @author luozhida
 * @create 2024-07-18
 */
public class SearchClient {


    final static int CLIENT_NUM = 15;
    final static int CLIENT_LOOP_NUM = 1;
    final static String COLLECTION_NAME = "xxxx";
    final static SimpleDateFormat SDF = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");

    final static Map<Integer, Pair<String, List<List<Float>>>> INNER_QUESTION_MAP = new HashMap<>();

    static {
        try {
            List<String> texts = FileUtils.readLines(new File(System.getProperty("user.dir") + "/src/test/resources/inner_embedding"), StandardCharsets.UTF_8);
            if (!texts.isEmpty()) {
                for (int i = 0; i < texts.size(); i++) {
                    String text = texts.get(i);
                    String[] array = text.split("#");
                    String embeddingText = array[1];
                    List<Float> floatList = new ArrayList<>();
                    for (String s : embeddingText.split(",")) {
                        floatList.add(Float.valueOf(s.trim()));
                    }
                    INNER_QUESTION_MAP.put(i, Pair.of(array[0], Arrays.asList(floatList)));
                }
            }
        } catch (Exception e) {
            throw new RuntimeException("读取内置问题向量文件失败", e);
        }
    }

    public static void main(String[] args) {
        System.out.println(SDF.format(new Date())+"开始查询");

        ComplexConnectConfig milvusConfig = new ComplexConnectConfig();
        milvusConfig.setHost("172.28.31.127");
        milvusConfig.setPort(19530);
        milvusConfig.setUser("root");
        milvusConfig.setPassword("Milvus");
        MilvusConnectionPool milvusConnectionPool = new MilvusConnectionPool("search-pool", milvusConfig, CLIENT_NUM);


        long startTimeMs = System.currentTimeMillis();
        List<CompletableFuture<Integer>> futures = new ArrayList<>(CLIENT_NUM);
        for (int i = 1; i <= CLIENT_NUM; i++) {
            futures.add(
                    CompletableFuture.supplyAsync(() -> {
                        MilvusServiceClient milvusServiceClient;
                        try {
                            milvusServiceClient = milvusConnectionPool.getConnection();
                        } catch (InterruptedException e) {
                            return 0;
                        }

                        int success = 0;
                        try {
                            for (int j = 1; j <= CLIENT_LOOP_NUM; j++) {
                                int randomKey = ThreadLocalRandom.current().nextInt(INNER_QUESTION_MAP.size());
                                Pair<String, List<List<Float>>> pair = INNER_QUESTION_MAP.get(randomKey);
                                try {
                                    R<SearchResults> resultsR = milvusServiceClient.search(SearchParam.newBuilder()
                                            .withCollectionName(COLLECTION_NAME)
                                            .withVectorFieldName("embedding")
                                            .withTopK(160)
                                            .withVectors(pair.getRight())
                                            .withOutFields(Arrays.asList("*"))
                                            .build());
                                    if (resultsR.getStatus() == R.Status.Success.getCode()) {
                                        System.out.println(SDF.format(new Date())+"查询成功,使用的问题下标:"
                                                +randomKey+",检索结果中最大分数为:" + resultsR.getData().getResults().getScores(0)
                                                +",检索出:"+resultsR.getData().getResults().getScoresCount()+"条");
                                        success++;
                                    }
                                } catch (Exception e) {
                                    System.out.println("查询出现异常:" + e.getMessage());
                                    e.printStackTrace();
                                }
                            }
                        } finally {
                            if (milvusServiceClient != null) {
                                milvusConnectionPool.releaseConnection(milvusServiceClient);
                            }
                        }
                        return success;
                    })
            );
        }
        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

        int allSuccess = 0;
        for (CompletableFuture<Integer> future : futures) {
            try {
                allSuccess = allSuccess + future.get();
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(e);
            }
        }
        long usedMill = System.currentTimeMillis() - startTimeMs;

        System.out.println(String.format("从【%s】查询成功共%s条记录,耗时%s毫秒,速率约%s条/秒",
                COLLECTION_NAME,
                allSuccess,
                usedMill,
                allSuccess <= 0 || usedMill <= 0 ? 0 : ((new BigDecimal(allSuccess))
                        .divide(new BigDecimal(
                                usedMill<=1000 ? 1 : TimeUnit.MILLISECONDS.toSeconds(usedMill)), 2, RoundingMode.HALF_DOWN))));
    }
}
